﻿#pragma once
#ifndef WsDispather_hpp
#define WsDispather_hpp

//#define UWS_NO_ZLIB

#include<functional>
#include<unordered_map>
#include<string>
#include<string_view>
#include<unordered_set>
#include"App.h"
#include"NetworkDef.hpp"
#include"../Concurrency/SpinLock.hpp"
#include"../Concurrency/Lock.hpp"
#include"../Malloc/MemoryPool.hpp"


namespace MMF
{

    template<typename UserDataT>
    class WsDispather
    {

        using WsMsgCallback = std::function<void(uWS::WebSocket<false,true,UserDataT>* ws, MMFNetwork::MMFMessageHeader* pkg)>;
        
        //内存池
        MMFAlloc::MemoryPool mMemPool;

        //消息与回调函数的映射
        //这里使用string_view时因为想减少从buffer到std::string的产生临时量内存申请
        std::unordered_map<std::string_view, WsMsgCallback> mMsgCallback;

        //互斥锁
        MMF::SpinLock mLock;

        //当消息未注册时调用的默认处理
        WsMsgCallback mDefaultCallback;
        WsMsgCallback mOpenCallback;
        WsMsgCallback mCloseCallback;

        //管理已经连接上来的套接字
        std::unordered_set<uWS::WebSocket<false,true,UserDataT>*> mWsScokets;

    public:

        WsDispather(const WsMsgCallback& defaultCallback)
         :mDefaultCallback(defaultCallback)
        {
            mOpenCallback = [this](auto* ws, auto* pkg){};
            mCloseCallback = [this](auto* ws, auto* pkg){};

            //添加默认处理的两个函数
            RegisterCallback("WsConnected", [this](auto* ws, auto* pkg){
                mWsScokets.insert(ws);
                mOpenCallback(ws, pkg);
            });
            RegisterCallback("WsClose", [this](auto* ws, auto* pkg){
                mWsScokets.erase(ws);
                mCloseCallback(ws, pkg);
            });
        }

        void RegisterClientOpenHanlder(const WsMsgCallback&& handler)
        {
            mLock.Lock();
            mOpenCallback = handler;
            mLock.Unlock();
        }

        void RegisterClientCloseHanlder(const WsMsgCallback&& handler)
        {
            mLock.Lock();
            mCloseCallback = handler;
            mLock.Unlock();
        }

        ~WsDispather(){}

        //处理消息
        void HandleMessage(uWS::WebSocket<false,true,UserDataT>* ws, void* pkg, size_t size)
        {
            //从内存池申请内存
            //MMFNetwork::MMFMessageHeader* msg = (MMFNetwork::MMFMessageHeader*)mMemPool.Alloc((uint32_t)size);
            MMFNetwork::MMFMessageHeader* msg = (MMFNetwork::MMFMessageHeader*)malloc((uint32_t)size);
            memcpy(msg, pkg, size);
            const char* msgName = msg->mName;

            //加锁
            mLock.Lock();

            //获取回调函数
            typename std::unordered_map<const std::string_view, WsMsgCallback>::iterator it = mMsgCallback.find(msgName);
            WsMsgCallback* handler = &mDefaultCallback;
            if (it != mMsgCallback.end())
            handler = &(it->second);

            //解锁
            mLock.Unlock();

            //调用回调
            (*handler)(ws, msg);

            //释放内存
            //mMemPool.Free(msg);
            free(msg);
        }

        //注册消息回调
        bool RegisterCallback(const char* msgName, WsMsgCallback callback)
        {
            MMF::LockGuard lk(mLock);

            typename std::unordered_map<const std::string_view, WsMsgCallback>::iterator it = mMsgCallback.find(msgName);
            if (it != mMsgCallback.end())
                return false;

            mMsgCallback.emplace(msgName, callback);

            return true;
        }

        //反注册消息回调
        bool UnregisterCallback(const char* msgName)
        {
            //加锁
            MMF::LockGuard lk(mLock);

            typename std::unordered_map<const std::string_view, WsMsgCallback>::iterator it = mMsgCallback.find(msgName);
            if (it == mMsgCallback.end())
                return false;

            mMsgCallback.erase(it);

            return true;
        }

        const std::unordered_set<uWS::WebSocket<false,true,UserDataT>*>& Session()
        {
            //加锁
            MMF::LockGuard lk(mLock);

            return mWsScokets;
        }
    };
}


#endif
